#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
高级代码混淆系统 - 实现源代码级别的混淆和保护
通过多种混淆技术使代码难以被理解和逆向工程
"""
import os
import sys
import re
import ast
import random
import string
import hashlib
import logging
import time
import inspect
from typing import Dict, List, Set, Tuple, Optional, Union, Any

# 导入项目模块
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))  # noqa
from utils import Logger
from utils.file_utils import FileUtils

class CodeObfuscator:
    """高级代码混淆器核心类"""
    def __init__(self):
        """初始化代码混淆器"""
        self.file_utils = FileUtils()
        
        # 混淆配置
        self.obfuscation_config = {
            "enable_obfuscation": True,
            "variable_renaming": True,
            "function_renaming": True,
            "class_renaming": True,
            "string_obfuscation": True,
            "control_flow_obfuscation": True,
            "dead_code_injection": False,
            "code_split": False,
            "encryption_level": "medium",  # low, medium, high
            "exclude_dirs": ["tests", "__pycache__", ".git", "docs"],
            "exclude_files": ["__init__.py", "setup.py"],
            "reserved_names": set(["__init__", "__main__", "__name__", "__file__", "__doc__"]) 
        }
        
        # 字符集用于生成随机名称
        self.characters = string.ascii_letters + string.digits
        
        Logger.log_info("高级代码混淆器初始化完成")
    
    def obfuscate_file(self, file_path: str, output_path: Optional[str] = None, config_overrides: Optional[Dict] = None) -> Tuple[bool, str, Optional[str]]:
        """混淆单个Python文件
        
        参数:
            file_path: 要混淆的文件路径
            output_path: 混淆后的输出文件路径，默认为原文件路径加上.obfuscated后缀
            config_overrides: 混淆配置覆盖
        
        返回:
            tuple: (是否成功, 消息, 混淆后的文件路径)
        """
        try:
            # 检查文件是否存在
            if not os.path.exists(file_path):
                error_msg = f"要混淆的文件不存在: {file_path}"
                Logger.log_error(error_msg)
                return False, error_msg, None
            
            # 检查文件类型
            _, ext = os.path.splitext(file_path)
            if ext.lower() != '.py':
                error_msg = f"不支持混淆的文件类型: {ext}"
                Logger.log_error(error_msg)
                return False, error_msg, None
            
            # 检查是否在排除列表中
            if self._is_excluded(file_path):
                Logger.log_info(f"文件在排除列表中，跳过混淆: {file_path}")
                return True, "文件在排除列表中，跳过混淆", None
            
            # 应用配置覆盖
            config = self.obfuscation_config.copy()
            if config_overrides:
                config.update(config_overrides)
            
            # 如果未启用混淆，直接返回成功
            if not config["enable_obfuscation"]:
                Logger.log_info(f"混淆功能未启用，跳过混淆: {file_path}")
                return True, "混淆功能未启用，跳过混淆", None
            
            # 确定输出路径
            if output_path is None:
                output_path = f"{file_path}.obfuscated"
            
            # 确保输出目录存在
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 读取文件内容
            with open(file_path, 'r', encoding='utf-8') as f:
                code_content = f.read()
            
            # 执行混淆
            obfuscated_code = self._obfuscate_code(code_content, config)
            
            # 保存混淆后的代码
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(obfuscated_code)
            
            Logger.log_info(f"文件混淆成功: {file_path} -> {output_path}")
            
            return True, "文件混淆成功", output_path
        except Exception as e:
            error_msg = f"混淆文件时发生错误: {str(e)}"
            Logger.log_error(error_msg)
            return False, error_msg, None
    
    def obfuscate_directory(self, dir_path: str, output_dir: Optional[str] = None, config_overrides: Optional[Dict] = None) -> Tuple[bool, str, Optional[Dict]]:
        """混淆目录中的所有Python文件
        
        参数:
            dir_path: 要混淆的目录路径
            output_dir: 混淆后的输出目录路径，默认为原目录路径加上.obfuscated后缀
            config_overrides: 混淆配置覆盖
        
        返回:
            tuple: (是否成功, 消息, 混淆统计信息)
        """
        try:
            # 检查目录是否存在
            if not os.path.exists(dir_path):
                error_msg = f"要混淆的目录不存在: {dir_path}"
                Logger.log_error(error_msg)
                return False, error_msg, None
            
            # 确定输出目录
            if output_dir is None:
                output_dir = f"{dir_path}.obfuscated"
            
            # 确保输出目录存在
            os.makedirs(output_dir, exist_ok=True)
            
            # 初始化统计信息
            stats = {
                "total_files": 0,
                "obfuscated_files": 0,
                "skipped_files": 0,
                "failed_files": 0,
                "start_time": time.time(),
                "end_time": None,
                "duration": None
            }
            
            # 递归遍历目录
            for root, dirs, files in os.walk(dir_path):
                # 过滤排除的目录
                dirs[:] = [d for d in dirs if d not in self.obfuscation_config["exclude_dirs"]]
                
                for file in files:
                    file_path = os.path.join(root, file)
                    stats["total_files"] += 1
                    
                    # 检查文件是否在排除列表中
                    if self._is_excluded(file_path):
                        stats["skipped_files"] += 1
                        continue
                    
                    # 检查文件类型
                    _, ext = os.path.splitext(file)
                    if ext.lower() != '.py':
                        stats["skipped_files"] += 1
                        continue
                    
                    # 计算相对路径
                    rel_path = os.path.relpath(file_path, dir_path)
                    output_path = os.path.join(output_dir, rel_path)
                    
                    # 确保输出目录存在
                    os.makedirs(os.path.dirname(output_path), exist_ok=True)
                    
                    # 混淆文件
                    success, message, _ = self.obfuscate_file(file_path, output_path, config_overrides)
                    
                    if success:
                        stats["obfuscated_files"] += 1
                    else:
                        stats["failed_files"] += 1
                        Logger.log_error(f"混淆文件失败: {file_path}, 错误: {message}")
            
            # 计算混淆持续时间
            stats["end_time"] = time.time()
            stats["duration"] = stats["end_time"] - stats["start_time"]
            
            Logger.log_info(
                f"目录混淆完成: {dir_path} -> {output_dir}, "
                f"总文件数: {stats['total_files']}, 混淆成功: {stats['obfuscated_files']}, "
                f"跳过: {stats['skipped_files']}, 失败: {stats['failed_files']}, "
                f"耗时: {stats['duration']:.2f}秒"
            )
            
            return True, "目录混淆完成", stats
        except Exception as e:
            error_msg = f"混淆目录时发生错误: {str(e)}"
            Logger.log_error(error_msg)
            return False, error_msg, None
    
    def _obfuscate_code(self, code_content: str, config: Dict) -> str:
        """执行代码混淆的核心逻辑"""
        try:
            # 解析代码为AST
            tree = ast.parse(code_content)
            
            # 创建混淆器访问者
            obfuscator_visitor = self._create_obfuscator_visitor(config)
            
            # 应用混淆
            modified_tree = obfuscator_visitor.visit(tree)
            
            # 将修改后的AST转换回代码
            obfuscated_code = ast.unparse(modified_tree) if hasattr(ast, 'unparse') else self._unparse_ast(modified_tree)
            
            # 应用字符串混淆（如果启用）
            if config["string_obfuscation"]:
                obfuscated_code = self._obfuscate_strings(obfuscated_code, config["encryption_level"])
            
            # 应用控制流混淆（如果启用）
            if config["control_flow_obfuscation"]:
                obfuscated_code = self._obfuscate_control_flow(obfuscated_code)
            
            # 注入死代码（如果启用）
            if config["dead_code_injection"]:
                obfuscated_code = self._inject_dead_code(obfuscated_code)
            
            # 分割代码（如果启用）
            if config["code_split"]:
                obfuscated_code = self._split_code(obfuscated_code)
            
            return obfuscated_code
        except Exception as e:
            Logger.log_error(f"代码混淆核心逻辑执行失败: {str(e)}")
            # 如果混淆失败，返回原始代码
            return code_content
    
    def _create_obfuscator_visitor(self, config: Dict) -> ast.NodeTransformer:
        """创建AST混淆器访问者"""
        # 变量名映射表
        var_name_map = {}
        # 函数名映射表
        func_name_map = {}
        # 类名映射表
        class_name_map = {}
        # 保留的名称
        reserved_names = config["reserved_names"]
        
        class ObfuscatorVisitor(ast.NodeTransformer):
            """AST混淆访问者"""
            def visit_FunctionDef(self, node):
                """访问函数定义"""
                if config["function_renaming"] and node.name not in reserved_names:
                    # 生成新的函数名
                    new_name = self._generate_random_name()
                    while new_name in func_name_map.values():
                        new_name = self._generate_random_name()
                    
                    func_name_map[node.name] = new_name
                    node.name = new_name
                
                # 继续访问函数体
                self.generic_visit(node)
                return node
            
            def visit_ClassDef(self, node):
                """访问类定义"""
                if config["class_renaming"] and node.name not in reserved_names:
                    # 生成新的类名
                    new_name = self._generate_random_name()
                    while new_name in class_name_map.values():
                        new_name = self._generate_random_name()
                    
                    class_name_map[node.name] = new_name
                    node.name = new_name
                
                # 继续访问类体
                self.generic_visit(node)
                return node
            
            def visit_Name(self, node):
                """访问变量名"""
                if isinstance(node.ctx, ast.Store) and config["variable_renaming"] and node.id not in reserved_names:
                    # 变量赋值时生成新名称
                    if node.id not in var_name_map:
                        new_name = self._generate_random_name()
                        while new_name in var_name_map.values() or new_name in func_name_map.values() or new_name in class_name_map.values():
                            new_name = self._generate_random_name()
                        
                        var_name_map[node.id] = new_name
                    
                    node.id = var_name_map[node.id]
                elif isinstance(node.ctx, ast.Load) and node.id in var_name_map:
                    # 变量引用时使用映射的新名称
                    node.id = var_name_map[node.id]
                elif isinstance(node.ctx, ast.Load) and node.id in func_name_map:
                    # 函数引用时使用映射的新名称
                    node.id = func_name_map[node.id]
                elif isinstance(node.ctx, ast.Load) and node.id in class_name_map:
                    # 类引用时使用映射的新名称
                    node.id = class_name_map[node.id]
                
                return node
            
            def _generate_random_name(self):
                """生成随机名称"""
                # 首字符不能是数字
                first_char = random.choice(string.ascii_letters)
                # 后续字符可以是字母和数字
                rest_chars = ''.join(random.choice(self.characters) for _ in range(random.randint(5, 15)))
                return first_char + rest_chars
        
        # 绑定外部方法到内部类
        ObfuscatorVisitor._generate_random_name = lambda self: ''.join(random.choice(self.characters) for _ in range(random.randint(8, 16)))
        
        return ObfuscatorVisitor()
    
    def _unparse_ast(self, tree: ast.AST) -> str:
        """将AST转换回代码字符串（兼容旧版Python）"""
        import astunparse
        return astunparse.unparse(tree)
    
    def _obfuscate_strings(self, code: str, encryption_level: str) -> str:
        """混淆代码中的字符串常量"""
        # 使用正则表达式匹配字符串
        def replace_string(match):
            string_value = match.group(0)
            # 根据加密级别选择不同的混淆方法
            if encryption_level == "high":
                # 高级加密：使用Base64编码和XOR加密
                encoded = base64.b64encode(string_value.encode('utf-8'))
                xored = bytes([b ^ 0x55 for b in encoded])
                return f"str(bytearray({list(xored)}).decode('utf-8','surrogateescape').encode('latin-1'))[2:-1]"
            elif encryption_level == "medium":
                # 中级加密：使用Base64编码
                encoded = base64.b64encode(string_value.encode('utf-8'))
                return f"__import__('base64').b64decode({encoded}).decode('utf-8')"
            else:
                # 低级加密：简单的字符反转
                reversed_str = string_value[::-1]
                return f"{reversed_str}[::-1]"
        
        # 替换所有字符串常量
        obfuscated_code = re.sub(r'"""[\s\S]*?"""|\'\'\'[\s\S]*?\'\'\'|"[^"]*"|\'[^\']*\'', replace_string, code)
        
        return obfuscated_code
    
    def _obfuscate_control_flow(self, code: str) -> str:
        """混淆代码的控制流程"""
        # 这是一个简化的实现，实际项目中可以使用更复杂的控制流混淆技术
        lines = code.split('\n')
        obfuscated_lines = []
        
        for line in lines:
            # 跳过空行和注释行
            if not line.strip() or line.strip().startswith('#'):
                obfuscated_lines.append(line)
                continue
            
            # 对条件语句进行简单混淆
            if 'if ' in line and ':' in line and not line.strip().startswith('elif') and not line.strip().startswith('else'):
                # 生成一个随机的逻辑表达式作为条件混淆
                random_expr = f"_x{random.randint(1000, 9999)} = {random.randint(0, 1)}; "
                obfuscated_lines.append(random_expr + line)
            else:
                obfuscated_lines.append(line)
        
        return '\n'.join(obfuscated_lines)
    
    def _inject_dead_code(self, code: str) -> str:
        """注入死代码"""
        lines = code.split('\n')
        obfuscated_lines = []
        
        for line in lines:
            # 添加原代码行
            obfuscated_lines.append(line)
            
            # 有一定概率在代码行后注入死代码
            if random.random() < 0.1 and line.strip() and not line.strip().startswith('#') and not line.strip().endswith(':'):
                # 生成死代码
                dead_code = f"    # 死代码注入 - {time.time()}\n    "
                dead_code += f"if {random.randint(0, 1)} == {random.randint(2, 10)}:  # 永远为False的条件\n"
                dead_code += f"        {random.choice(['pass', 'print(None)', 'x = None'])}  # 无意义的操作\n"
                
                # 缩进死代码以匹配上下文
                indent = len(line) - len(line.lstrip())
                dead_code = ' ' * indent + dead_code.replace('    ', ' ' * (indent + 4))
                
                obfuscated_lines.append(dead_code)
        
        return '\n'.join(obfuscated_lines)
    
    def _split_code(self, code: str) -> str:
        """分割代码以增加复杂性"""
        # 这是一个简化的实现，实际项目中可以使用更复杂的代码分割技术
        # 此方法仅在高级混淆模式下启用
        return code
    
    def _is_excluded(self, file_path: str) -> bool:
        """检查文件是否在排除列表中"""
        # 检查文件是否在排除文件列表中
        file_name = os.path.basename(file_path)
        if file_name in self.obfuscation_config["exclude_files"]:
            return True
        
        # 检查文件所在目录是否在排除目录列表中
        for dir_name in self.obfuscation_config["exclude_dirs"]:
            if dir_name in file_path.split(os.path.sep):
                return True
        
        return False
    
    def generate_obfuscation_map(self, input_file: str, output_file: str) -> bool:
        """生成混淆映射文件，用于调试和开发"""
        try:
            # 此方法仅在调试模式下有效
            # 实际项目中可以实现此方法以生成变量名、函数名、类名的映射关系
            Logger.log_info(f"生成混淆映射文件: {input_file} -> {output_file}")
            return True
        except Exception as e:
            Logger.log_error(f"生成混淆映射文件失败: {str(e)}")
            return False

# 导入必要的模块
import base64

if __name__ == "__main__":
    # 示例用法
    obfuscator = CodeObfuscator()
    
    # 混淆单个文件
    success, message, output_path = obfuscator.obfuscate_file(
        "example.py",
        config_overrides={"encryption_level": "high"}
    )
    print(f"混淆结果: {success}, 消息: {message}, 输出路径: {output_path}")
    
    # 混淆整个目录
    # success, message, stats = obfuscator.obfuscate_directory("src")
    # print(f"目录混淆结果: {success}, 消息: {message}, 统计: {stats}")